package tree

import (
	"data-structure/queue"
	"fmt"
	"math"
)

type AVLTree struct {
	Head *AVLNode
}

type AVLNode struct {
	V    int
	L, R *AVLNode
	H, C int
}

//construct binary search tree

//其实就是递归选取根节点的过程
func InitAVL(input []int) *AVLTree {
	fsort(input)
	return &AVLTree{
		Head: CompleteBinarySeachTree(input),
	}
}

//层序遍历, 使用一个队列，从队头取出元素，把该元素的子树放到队列末尾。
//所以需要知道这一层的节点数量，即len
func LevelOrder(root *AVLNode) {
	nodeQueue := queue.NewNodeQueue()
	nodeQueue.Append(root)
	levelOrderHelper(nodeQueue)
}

func levelOrderHelper(nodeQueue *queue.NodeQueue) {
	length := nodeQueue.Length()
	if length == 0 {
		return
	}
	for i := 0; i < length; i++ {
		ele, _ := nodeQueue.Pop()
		eleNode := ele.(*AVLNode)
		fmt.Print("V:", eleNode.V, "H:", eleNode.H, "    ")
		if eleNode.L != nil {
			nodeQueue.Append(eleNode.L)
		}
		if eleNode.R != nil {
			nodeQueue.Append(eleNode.R)
		}
	}
	fmt.Println()
	levelOrderHelper(nodeQueue)
}

func fsort(input []int) {
	if len(input) <= 1 {
		return
	}
	mid := input[0]
	left, right := 1, len(input)-1
	for left <= right {
		for ; left <= right && input[left] < mid; left++ {
		}
		for ; left <= right && input[right] >= mid; right-- {
		}
		//此时left和right必定不相等，
		if left < right {
			input[left], input[right] = input[right], input[left]
			left++
			right--
		}
	}
	input[right], input[0] = input[0], input[right]
	//right min = 0
	//left max = len(input)
	fsort(input[:right])
	fsort(input[left:])
}

func (avl *AVLTree) Insert(value int) {
	// 1. 如果还没有元素，给根节点赋值；
	if avl.Head == nil {
		avl.Head = &AVLNode{
			V: value,
			C: 1,
			H: 0,
		}
	}
	//2. 如果有元素，找到要插入的位置；并存储查找的路径；
	avlNode, alreadyExists, path := avl.Find(value)
	//3. 插入
	if alreadyExists {
		avlNode.C++
	} else {
		newNodePtr := &AVLNode{
			V: value,
			H: 0,
			C: 1,
		}
		if avlNode.V > value {
			avlNode.L = newNodePtr
		}
		if avlNode.V < value {
			avlNode.R = newNodePtr
		}
		*path = append(*path, newNodePtr)
	}
	root := &AVLNode{R: avl.Head}
	for k := len(*path) - 1; k >= 0; k-- {

		eleNode := (*path)[k]
		//4.更新高度，高度变化只和查找经过路径上的节点有关。
		eleNode.updateHeight()
		//5. 查看当前节点是否平衡
		if !eleNode.isBalance() {
			if k != 0 {
				newCRoot := justify((*path)[k:])
				if (*path)[k-1].V > newCRoot.V {
					(*path)[k-1].L = newCRoot
				} else {
					(*path)[k-1].R = newCRoot
				}
			} else {
				newCRoot := justify((*path)[k:])
				root.R = newCRoot
			}
		}
	}
	avl.Head = root.R
}

func max(arg1, arg2 int) int {
	if arg1 > arg2 {
		return arg1
	}
	return arg2
}

func (avl *AVLTree) Find(value int) (*AVLNode, bool, *[]*AVLNode) {
	path := []*AVLNode{}
	return find(avl.Head, value, &path)
}

func find(avlNode *AVLNode, value int, path *[]*AVLNode) (*AVLNode, bool, *[]*AVLNode) {
	var nextAvlNode *AVLNode
	if avlNode.V == value {
		return avlNode, true, path
	} else if avlNode.V > value {
		if avlNode.L == nil {
			*path = append(*path, avlNode)
			return avlNode, false, path
		}
		nextAvlNode = avlNode.L
	} else if avlNode.V < value {
		if avlNode.R == nil {
			*path = append(*path, avlNode)
			return avlNode, false, path
		}
		nextAvlNode = avlNode.R
	}
	*path = append(*path, avlNode)
	return find(nextAvlNode, value, path)
}

//把父节点也传入进来, 旋转之后,将父节点的子节点替换
func justify(input []*AVLNode) *AVLNode {
	//通过大小顺序判断选择什么旋转,如果要不平衡, 节点路径需要往回倒至少3个节点.
	//每次调整都传入路径上的三个节点.
	//RR
	if input[0].V > input[1].V && input[1].V > input[2].V {
		input[0] = input[0].RightRotate()
	} else if input[0].V < input[1].V && input[1].V < input[2].V {
		input[0] = input[0].LeftRotate()
	} else if input[0].V > input[1].V && input[1].V < input[2].V {
		input[0] = input[0].LRDoubleRotate()
	} else {
		input[0] = input[0].RLDoubleRotate()
	}
	return input[0]
}

//旋转之后返回根节点,
/**	  1                      2
        \                   / \
	      2        =>      1   3
            \
             3
*/
func (avl *AVLNode) LeftRotate() *AVLNode {
	root := avl.R
	avl.R = avl.R.L
	root.L = avl
	root.L.updateHeight()
	root.updateHeight()
	return root
}

/**				3                    2
               / \                  / \
	          2   a    =>          1   3
             / \                  / \  /\
            1   b                c   d b a
           / \
          c   d
*/
func (avl *AVLNode) RightRotate() *AVLNode {
	root := avl.L
	avl.L = avl.L.R
	root.R = avl
	root.R.updateHeight()
	root.updateHeight()
	return root
}

func (avl *AVLNode) LRDoubleRotate() *AVLNode {
	root := avl.L.R
	tempL := root.L
	tempR := root.R
	root.L = avl.L
	root.R = avl
	avl.L.R = tempL
	avl.L = tempR
	root.L.updateHeight()
	root.R.updateHeight()
	root.updateHeight()
	return root
}

func (avl *AVLNode) RLDoubleRotate() *AVLNode {
	root := avl.R.L
	tempL := root.L
	tempR := root.R
	root.L = avl
	root.R = avl.R
	avl.R = tempL
	avl.R.L = tempR
	root.L.updateHeight()
	root.R.updateHeight()
	root.updateHeight()
	return root
}

func (avlNode *AVLNode) isBalance() bool {
	if math.Abs(float64(height(avlNode.L)-height(avlNode.R))) <= 1 {
		return true
	}
	return false
}

func height(avlNode *AVLNode) int {
	if avlNode == nil {
		return -1
	}
	return avlNode.H
}

func (avlNode *AVLNode) updateHeight() {
	avlNode.H = max(height(avlNode.R), height(avlNode.L)) + 1
}

//@TODO: 节点删除
